from psycopg2 import sql

class DataService:
    """All database operations go this way."""
    def __init__(self, conn):
        self.conn = conn

    # ---------- universal method ----------
    def get_columns(self, table: str):
        """Returns whether [(column name, data type, self-increment)] is self-increment."""
        return self._safe_execute("""
            SELECT column_name, data_type,
                   CASE WHEN column_default LIKE 'nextval%%' THEN 1 ELSE 0 END AS is_serial
            FROM information_schema.columns
            WHERE table_schema=%s AND table_name=%s
            ORDER BY ordinal_position
        """, table.replace('"', '').split('.'))

    def get_primary_keys(self, table: str):
        """Returns a list of primary key column names."""
        schema, tbl = table.replace('"', '').split('.')
        rows = self._safe_execute("""
            SELECT kcu.column_name
            FROM information_schema.table_constraints tc
            JOIN information_schema.key_column_usage kcu
              ON tc.constraint_name = kcu.constraint_name
             AND tc.table_schema = kcu.table_schema
            WHERE tc.constraint_type = 'PRIMARY KEY'
              AND tc.table_schema=%s AND tc.table_name=%s
        """, (schema, tbl))
        return [row[0] for row in rows]

    # ---------- Query ----------
    def query(self, table: str, where_dict: dict):
        """
        Where_dict: {column name: value}
        Returns (list of column names, list of data rows)
        """
        if not where_dict:
            where_sql = sql.SQL("")
            params = []
        else:
            conditions = [sql.SQL("{} = {}").format(sql.Identifier(k), sql.Placeholder()) for k in where_dict]
            where_sql = sql.SQL(" WHERE ") + sql.SQL(" AND ").join(conditions)
            params = list(where_dict.values())

        query_sql = sql.SQL("SELECT * FROM {}").format(
            sql.Identifier(*table.replace('"', '').split('.'))) + where_sql

        try:
            with self.conn.cursor() as cur:
                cur.execute(query_sql, params)
                return [desc[0] for desc in cur.description], cur.fetchall()
        except Exception:
            self.conn.rollback()
            raise

    # ---------- Insert ----------
    def insert(self, table: str, value_dict: dict):
        """
        Value_dict: {column name: value}
        returns the number of rows affected.
        """
        cols = list(value_dict.keys())
        vals = list(value_dict.values())
        for c in cols.copy():
            if value_dict[c] == '':
                idx = cols.index(c)
                cols.pop(idx)
                vals.pop(idx)
        insert_sql = sql.SQL("INSERT INTO {} ({}) VALUES ({})").format(
            sql.Identifier(*table.replace('"', '').split('.')),
            sql.SQL(", ").join(map(sql.Identifier, cols)),
            sql.SQL(", ").join(sql.Placeholder() * len(cols))
        )
        try:
            with self.conn.cursor() as cur:
                cur.execute(insert_sql, vals)
                self.conn.commit()
                return cur.rowcount
        except Exception:
            self.conn.rollback()
            raise

    # ---------- Update ----------
    def update(self, table: str, value_dict: dict, pk_dict: dict):
        """
        value_dict: Column to be changed
        pk_dict:    Primary Key Column = Value Locating Row
        """
        set_part = sql.SQL(", ").join(
            sql.SQL("{} = {}").format(sql.Identifier(k), sql.Placeholder()) for k in value_dict
        )
        where_part = sql.SQL(" AND ").join(
            sql.SQL("{} = {}").format(sql.Identifier(k), sql.Placeholder()) for k in pk_dict
        )
        update_sql = sql.SQL("UPDATE {} SET {} WHERE {}").format(
            sql.Identifier(*table.replace('"', '').split('.')), set_part, where_part
        )
        try:
            with self.conn.cursor() as cur:
                cur.execute(update_sql, list(value_dict.values()) + list(pk_dict.values()))
                self.conn.commit()
                return cur.rowcount
        except Exception:
            self.conn.rollback()
            raise

    # ---------- delete ----------
    def delete(self, table: str, pk_dict: dict):
        where_part = sql.SQL(" AND ").join(
            sql.SQL("{} = {}").format(sql.Identifier(k), sql.Placeholder()) for k in pk_dict
        )
        delete_sql = sql.SQL("DELETE FROM {} WHERE {}").format(
            sql.Identifier(*table.replace('"', '').split('.')), where_part
        )
        try:
            with self.conn.cursor() as cur:
                cur.execute(delete_sql, list(pk_dict.values()))
                self.conn.commit()
                return cur.rowcount
        except Exception:
            self.conn.rollback()
            raise

    # ---------- helper ----------
    def _safe_execute(self, sql_str: str, params: tuple | list):
        """Execute the SELECT class statement, and if it is abnormal, rollback first and then throw it."""
        try:
            with self.conn.cursor() as cur:
                cur.execute(sql_str, params)
                return cur.fetchall()
        except Exception:
            self.conn.rollback()
            raise

    # ---------- complex query ----------
    def execute_complex_query(self, query_name: str, params: dict = None):
        """Execute predefined complex queries"""
        if params is None:
            params = {}

        queries = {
            # Query 1: Calculate monthly income of departments (using function)
            "dept_monthly_revenue": """
                SELECT calculate_dept_monthly_revenue(%(dept_id)s, %(year)s, %(month)s) AS monthly_revenue
            """,

            # Query 2: Obtain the complete medical history of patients (using the function)
            "patient_medical_history": """
                SELECT * from get_patient_medical_history(%(patient_id)s)
            """,

            # Query 3: High-frequency query-Search patients at the front desk (using index)
            "search_patients": """
                SELECT patient_id, first_name, last_name, phone, email
                FROM patients 
                WHERE last_name ILIKE %(name_pattern)s OR first_name ILIKE %(name_pattern)s
                ORDER BY last_name, first_name
                LIMIT 50
            """
        }

        if query_name not in queries:
            raise ValueError(f"Unknown query: {query_name}")

        sql = queries[query_name]
        return self._safe_execute_with_columns(sql, params)

    def _safe_execute_with_columns(self, sql_str: str, params: dict):
        """Executes the query and returns column names and results"""
        try:
            with self.conn.cursor() as cur:
                cur.execute(sql_str, params)
                if cur.description:
                    columns = [desc[0] for desc in cur.description]
                    rows = cur.fetchall()
                    return columns, rows
                else:
                    return [], []
        except Exception:
            self.conn.rollback()
            raise